//
// Created by Admin on 2024/1/3.
//

#include "RtpParser.h"

#if defined(_WIN32)
#include <winsock2.h>
#else
#include <arpa/inet.h>
#endif

#include <sstream>
#include <string>
#include <iostream>
#include <fstream>
using namespace std;

#include "rtprawpacket.h"
#include "rtppacket.h"
using namespace jrtplib;

void SaveRealData(const char* fileName, const char* pszData, int nSize)
{
    char sModuleName[MAX_PATH] = { 0 };
    ::GetModuleFileNameA(NULL, sModuleName, MAX_PATH);
    string app_dir(sModuleName);
    string::size_type nLastBacklash = app_dir.rfind('\\', app_dir.size());
    if (nLastBacklash != string::npos)
        app_dir.erase(nLastBacklash + 1);
    string dir = app_dir.append(fileName);

    ofstream os(dir, ios::binary | ios::app);
    if (os.is_open())
    {
        os.write(pszData, nSize);
        os.close();
    }
}

//协议解析最大缓存4兆数据
static constexpr size_t kMaxCacheSize = 4 * 1024 * 1024;

//===========================
// \brief RtpHeader
//===========================
#define AV_RB16(x) ((((const uint8_t *)(x))[0] << 8) | ((const uint8_t *)(x))[1])

size_t RtpHeader::getCsrcSize() const {
    return csrc << 2;
}

uint8_t *RtpHeader::getCsrcData() {
    if (!csrc) {
        return nullptr;
    }
    return &payload;
}

size_t RtpHeader::getExtSize() const {
    // rtp有ext
    if (!ext) {
        return 0;
    }
    auto ext_ptr = &payload + getCsrcSize();
    // uint16_t reserved = AV_RB16(ext_ptr);
    // 每个ext占用4字节
    return AV_RB16(ext_ptr + 2) << 2;
}

uint16_t RtpHeader::getExtReserved() const {
    // rtp有ext
    if (!ext) {
        return 0;
    }
    auto ext_ptr = &payload + getCsrcSize();
    return AV_RB16(ext_ptr);
}

uint8_t *RtpHeader::getExtData() {
    if (!ext) {
        return nullptr;
    }
    auto ext_ptr = &payload + getCsrcSize();
    // 多出的4个字节分别为reserved、ext_len
    return ext_ptr + 4;
}

uint8_t *RtpHeader::getPayloadData() {
    return &payload + getPayloadOffset();
}

ssize_t RtpHeader::getPayloadSize(size_t rtp_size) const {
    auto invalid_size = getPayloadOffset() + getPaddingSize(rtp_size);
    return (ssize_t)rtp_size - invalid_size - kRtpHeaderSize;
}

std::string RtpHeader::dumpString(size_t rtp_size) const {
    std::stringstream ss;
    ss << "version:" << (int)version << "\r\n";
    ss << "padding:" << getPaddingSize(rtp_size) << "\r\n";
    ss << "ext:" << getExtSize() << "\r\n";
    ss << "csrc:" << getCsrcSize() << "\r\n";
    ss << "mark:" << (int)mark << "\r\n";
    ss << "pt:" << (int)pt << "\r\n";
    ss << "seq:" << ntohs(seq) << "\r\n";
    ss << "stamp:" << ntohl(stamp) << "\r\n";
    ss << "ssrc:" << ntohl(ssrc) << "\r\n";
    ss << "rtp size:" << rtp_size << "\r\n";
    ss << "payload offset:" << getPayloadOffset() << "\r\n";
    ss << "payload size:" << getPayloadSize(rtp_size) << "\r\n";
    return ss.str();
}

size_t RtpHeader::getPayloadOffset() const {
    // 有ext时，还需要忽略reserved、ext_len 4个字节
    return getCsrcSize() + (ext ? (4 + getExtSize()) : 0);
}

size_t RtpHeader::getPaddingSize(size_t rtp_size) const {
    if (!padding) {
        return 0;
    }
    auto end = (uint8_t *)this + rtp_size - 1;
    return *end;
}

//===========================
// \brief RtpHeader
//===========================
RtpBuffer::RtpBuffer() {
    _erase_head = 0;
    _erase_tail = 0;
}

RtpBuffer::RtpBuffer(std::string str) {
    _str = std::move(str);
    _erase_head = 0;
    _erase_tail = 0;
}

RtpBuffer& RtpBuffer::operator=(std::string str) {
    _str = std::move(str);
    _erase_head = 0;
    _erase_tail = 0;
    return *this;
}

RtpBuffer::RtpBuffer(const char* str) {
    _str = str;
    _erase_head = 0;
    _erase_tail = 0;
}

RtpBuffer& RtpBuffer::operator=(const char* str) {
    _str = str;
    _erase_head = 0;
    _erase_tail = 0;
    return *this;
}

RtpBuffer::RtpBuffer(RtpBuffer&& that) {
    _str = std::move(that._str);
    _erase_head = that._erase_head;
    _erase_tail = that._erase_tail;
    that._erase_head = 0;
    that._erase_tail = 0;
}

RtpBuffer& RtpBuffer::operator=(RtpBuffer&& that) {
    _str = std::move(that._str);
    _erase_head = that._erase_head;
    _erase_tail = that._erase_tail;
    that._erase_head = 0;
    that._erase_tail = 0;
    return *this;
}

RtpBuffer::RtpBuffer(const RtpBuffer& that) {
    _str = that._str;
    _erase_head = that._erase_head;
    _erase_tail = that._erase_tail;
}

RtpBuffer& RtpBuffer::operator=(const RtpBuffer& that) {
    _str = that._str;
    _erase_head = that._erase_head;
    _erase_tail = that._erase_tail;
    return *this;
}

char* RtpBuffer::data() const {
    return (char*)_str.data() + _erase_head;
}

size_t RtpBuffer::size() const {
    return _str.size() - _erase_tail - _erase_head;
}

RtpBuffer& RtpBuffer::erase(size_t pos, size_t n) {
    if (pos == 0) {
        //移除前面的数据
        if (n != std::string::npos) {
            //移除部分
            if (n > size()) {
                //移除太多数据了
                throw std::out_of_range("RtpBuffer::erase out_of_range in head");
            }
            //设置起始便宜量
            _erase_head += n;
            data()[size()] = '\0';
            return *this;
        }
        //移除全部数据
        _erase_head = 0;
        _erase_tail = _str.size();
        data()[0] = '\0';
        return *this;
    }

    if (n == std::string::npos || pos + n >= size()) {
        //移除末尾所有数据
        if (pos >= size()) {
            //移除太多数据
            throw std::out_of_range("RtpBuffer::erase out_of_range in tail");
        }
        _erase_tail += size() - pos;
        data()[size()] = '\0';
        return *this;
    }

    //移除中间的
    if (pos + n > size()) {
        //超过长度限制
        throw std::out_of_range("RtpBuffer::erase out_of_range in middle");
    }
    _str.erase(_erase_head + pos, n);
    return *this;
}

RtpBuffer& RtpBuffer::append(const RtpBuffer& str) {
    return append(str.data(), str.size());
}

RtpBuffer& RtpBuffer::append(const std::string& str) {
    return append(str.data(), str.size());
}

RtpBuffer& RtpBuffer::append(const char* data) {
    return append(data, strlen(data));
}

RtpBuffer& RtpBuffer::append(const char* data, size_t len) {
    if (len <= 0) {
        return *this;
    }
    if (_erase_head > _str.capacity() / 2) {
        moveData();
    }
    if (_erase_tail == 0) {
        _str.append(data, len);
        return *this;
    }
    _str.insert(_erase_head + size(), data, len);
    return *this;
}

void RtpBuffer::push_back(char c) {
    if (_erase_tail == 0) {
        _str.push_back(c);
        return;
    }
    data()[size()] = c;
    --_erase_tail;
    data()[size()] = '\0';
}

RtpBuffer& RtpBuffer::insert(size_t pos, const char* s, size_t n) {
    _str.insert(_erase_head + pos, s, n);
    return *this;
}

RtpBuffer& RtpBuffer::assign(const char* data) {
    return assign(data, strlen(data));
}

RtpBuffer& RtpBuffer::assign(const char* data, size_t len) {
    if (len <= 0) {
        return *this;
    }
    if (data >= _str.data() && data < _str.data() + _str.size()) {
        _erase_head = data - _str.data();
        if (data + len > _str.data() + _str.size()) {
            throw std::out_of_range("RtpBuffer::assign out_of_range");
        }
        _erase_tail = _str.data() + _str.size() - (data + len);
        return *this;
    }
    _str.assign(data, len);
    _erase_head = 0;
    _erase_tail = 0;
    return *this;
}

void RtpBuffer::clear() {
    _erase_head = 0;
    _erase_tail = 0;
    _str.clear();
}

char& RtpBuffer::operator[](size_t pos) {
    if (pos >= size()) {
        throw std::out_of_range("RtpBuffer::operator[] out_of_range");
    }
    return data()[pos];
}

const char& RtpBuffer::operator[](size_t pos) const {
    return (*const_cast<RtpBuffer*>(this))[pos];
}

size_t RtpBuffer::capacity() const {
    return _str.capacity();
}

void RtpBuffer::reserve(size_t size) {
    _str.reserve(size);
}

void RtpBuffer::resize(size_t size, char c) {
    _str.resize(size, c);
    _erase_head = 0;
    _erase_tail = 0;
}

bool RtpBuffer::empty() const {
    return size() <= 0;
}

std::string RtpBuffer::substr(size_t pos, size_t n) const {
    if (n == std::string::npos) {
        //获取末尾所有的
        if (pos >= size()) {
            throw std::out_of_range("RtpBuffer::substr out_of_range");
        }
        return _str.substr(_erase_head + pos, size() - pos);
    }

    //获取部分
    if (pos + n > size()) {
        throw std::out_of_range("RtpBuffer::substr out_of_range");
    }
    return _str.substr(_erase_head + pos, n);
}

void RtpBuffer::moveData() {
    if (_erase_head) {
        _str.erase(0, _erase_head);
        _erase_head = 0;
    }
}

//===========================
// \brief RtpSpliter
//===========================
RtpSpliter::RtpSpliter() {
    setMaxCacheSize(0);
}

void RtpSpliter::input(const char *data, size_t len) {
    {
        auto size = remainDataSize();
        if (size > _max_cache_size) {
            //缓存太多数据无法处理，重置缓存
            reset();
        }
    }

    const char* ptr = data;
    if (!_remain_data.empty()) {
        _remain_data.append(data, len);
        data = ptr = _remain_data.data();
        len = _remain_data.size();
    }

    //确保ptr最后一个字节为0，防止strstr越界
    char& tail_ref = ((char*)ptr)[len];
    char tail_tmp = tail_ref;
    tail_ref = 0;

    //数据按请求头处理
    const char* index = nullptr;
    _remain_data_size = len;
    while (_content_len == 0 && _remain_data_size > 0 && (index = onSearchPacketTail(ptr, _remain_data_size)) != nullptr) {
        if (index == ptr)
            break;

        if (index < ptr || index > ptr + _remain_data_size) {
            throw std::out_of_range("rtp split out of range");
        }

        const char* header_ptr = ptr;
        ssize_t header_size = index - ptr;
        ptr = index;
        _remain_data_size = len - (ptr - data);
        _content_len = onRecvHeader(header_ptr, header_size);
    }

    if (_remain_data_size <= 0) {
        //没有数据，清空缓存
        _remain_data.clear();
        return;
    }

    if (_content_len < 0)
    {
        //无效数据，清空缓存
        _remain_data.clear();
        _content_len = 0;
        return;
    }

    //恢复末尾字节
    tail_ref = tail_tmp;

    if (_content_len == 0)
    {
        //处理剩余数据
        _remain_data.assign(ptr, _remain_data_size);
        return;
    }
}

void RtpSpliter::reset() {
    _content_len = 0;
    _remain_data_size = 0;
    _remain_data.clear();
}

size_t RtpSpliter::remainDataSize() {
    return _remain_data_size;
}

const char *RtpSpliter::remainData() const {
    return _remain_data.data();
}

void RtpSpliter::setMaxCacheSize(size_t max_cache_size) {
    if (!_max_cache_size) {
        max_cache_size = kMaxCacheSize;
    }
    _max_cache_size = max_cache_size;
}

ssize_t RtpSpliter::onRecvHeader(const char *data, size_t len) {
    //忽略偏移量
    data += _offset;
    len -= _offset;
    return onRtpPacket(data, len);
}

void RtpSpliter::onRecvContent(const char *data, size_t len) {

}

const char *RtpSpliter::onSearchPacketTail(const char *data, size_t len) {
    if (len < 4) {
        //数据不够
        return nullptr;
    }

    //两个字节的rtp头
    _offset = 2;

    //这是rtp包
    uint16_t length = (((uint8_t *) data)[0] << 8) | ((uint8_t *) data)[1];
    if (len < (size_t)(length + 2)) {
        //数据不够
        return nullptr;
    }
    //返回rtp包末尾
    return data + 2 + length;
}

void RtpSpliter::setContentLen(ssize_t content_len) {
    _content_len = content_len;
}

bool RtpSpliter::isRtp(const char* data, size_t len)
{
    if (len < 2)
        return false;
    RtpHeader* header = (RtpHeader*)data;
    return ((header->pt < 64) || (header->pt >= 96)) && header->version == kRtpVersion;
}

//===========================
// \brief RtpParser
//===========================
RtpParser::RtpParser(bool is_udp)
    : _is_udp(is_udp) {
}

RtpParser::~RtpParser() {

}

void RtpParser::input(const char *data, size_t len) {
    if (_is_udp) {
        onRtpPacket(data, len);
        return;
    }
    RtpSpliter::input(data, len);
}

ssize_t RtpParser::onRtpPacket(const char *data, size_t len) {
    if (isRtp(data, len))
    {
        jrtplib::RTPTime objRTPTime(0);
        jrtplib::RTPRawPacket objRTPRawPacket((uint8_t*)data, len, NULL, objRTPTime, true, nullptr);

        RTPPacket objRTPPacket(objRTPRawPacket, true);
        objRTPRawPacket.ZeroData();

        if (objRTPPacket.GetCreationError() >= 0)
        {
            int nFrameType = 0;
            if (objRTPPacket.GetExtensionData())
            {
                memcpy(&nFrameType, objRTPPacket.GetExtensionData(), sizeof(int));
            }

            SaveRealData("pack111", (const char*)objRTPPacket.GetPayloadData(), objRTPPacket.GetPayloadLength());
        }
        return 0;
    }
    return -1;
}



